// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: xguo

#include "cc/coin2/base/scheduler.h"

#include <random>
#include <utility>

#include "coin2/base/log.h"
#include "coin2/base/time.h"

Scheduler::Scheduler(boost::asio::io_context* ioc) : ioc_(ioc), strand_(nullptr) {
  LOG(INFO) << "[Scheduler] check period: " << check_period_ns_ << ", " << this;
  task_vec_.reserve(1024);
}

Scheduler::Scheduler(boost::asio::io_context::strand* strand) : ioc_(nullptr), strand_(strand) {
  LOG(INFO) << "[Scheduler] check period: " << check_period_ns_ << ", " << this;
  task_vec_.reserve(1024);
}

void Scheduler::SetCheckPeriod(int64_t check_period_ns) {
  check_period_ns_ = check_period_ns;
  LOG(INFO) << "[Scheduler] check period set to: " << check_period_ns_;
}

void Scheduler::Stop() { running_ = false; }

void Scheduler::AddPeriodicTask(
    Task::TaskFunc task_func,
    int64_t period,
    const std::string& task_name) {
  auto task = std::make_unique<Task>(task_func, period, GetCurrentTimestamp(), task_name);
  task_queue_.push(task.get());
  tasks_.push_back(std::move(task));
}

void Scheduler::AddPeriodicTask(
    Task::TaskFunc task_func,
    int64_t period,
    const std::string& task_name,
    int64_t delay) {
  auto task = std::make_unique<Task>(task_func, period, GetCurrentTimestamp() + delay, task_name);
  task_queue_.push(task.get());
  tasks_.push_back(std::move(task));
}

void Scheduler::CallLater(Task::TaskFunc task_func, int64_t delay, const std::string& task_name) {
  auto task = std::make_unique<Task>(task_func, -1, GetCurrentTimestamp() + delay, task_name);
  task_queue_.push(task.get());
  tasks_.push_back(std::move(task));
}

void Scheduler::StartImpl(boost::asio::yield_context yield_context) {
  const auto check_period = boost::posix_time::milliseconds(check_period_ns_ / 1000'000);
  boost::asio::io_context* ioc = nullptr;
  if (ioc_ != nullptr) {
    ioc = ioc_;
  } else if (strand_ != nullptr) {
    ioc = &(strand_->context());
  } else {
    THROW() << "timer has no context!";
  }

  boost::asio::deadline_timer timer(*ioc);
  running_ = true;
  std::default_random_engine generator;
  generator.seed(time(nullptr));
  std::uniform_int_distribution<int> distribution(0, 1000'000);
  auto randint = std::bind(distribution, generator);
  while (running_) {
    const auto curr_time = GetCurrentTimestamp();
    while (!task_queue_.empty()) {
      auto* task = task_queue_.top();
      if (task->next_run_time <= curr_time) {
        task_queue_.pop();
        task_vec_.push_back(task);
        DCHECK(reinterpret_cast<int64_t>(task) > 1000);
      } else {
        break;
      }
    }

    if (task_vec_.empty()) {
      timer.expires_from_now(check_period);
      timer.async_wait(yield_context);
      continue;
    }

    // Randomly pick one task.
    auto idx = randint() % task_vec_.size();
    auto task = task_vec_.at(idx);
    try {
      task->Run();
    } catch (std::exception& e) {
      LOG(ERROR) << "Task failed! task_name: " << task->task_name;
      LOG(ERROR) << "task exception: " << e.what();
    }

    if (task->period > 0) {
      task->ScheduleNextRun(curr_time);
      task_queue_.push(task);
    }
    if (task_vec_.size() > 10) {
      LOG_EVERY_NS(WARNING, 10'000'000'000, curr_time)
          << "it's too busy to process tasks! num of pending: " << task_vec_.size() - 1
          << ", num of schedule: " << task_queue_.size() + task_vec_.size() - 1;
    }

    for (size_t n = 0; n < task_vec_.size(); n++) {
      if (n != idx) {
        task_queue_.push(task_vec_[n]);
      }
    }
    task_vec_.clear();
    timer.expires_from_now(check_period);
    timer.async_wait(yield_context);
  }

  while (!task_queue_.empty()) {
    task_queue_.pop();
  }
  tasks_.clear();
  LOG(INFO) << "scheduler stopped! " << this;
}

void Scheduler::Start() {
  std::weak_ptr<Scheduler> self(shared_from_this());
  if (strand_ != nullptr) {
    LOG(INFO) << "scheduler started! " << this << "," << strand_;
    boost::asio::spawn(*strand_, [self](boost::asio::yield_context yc) {
      if (auto s = self.lock()) {
        s->StartImpl(std::move(yc));
      } else {
        LOG(ERROR) << "scheduler is already destroyed";
      }
    });
  } else if (ioc_ != nullptr) {
    LOG(INFO) << "scheduler started! " << this << "," << ioc_;
    boost::asio::spawn(*ioc_, [self](boost::asio::yield_context yc) {
      if (auto s = self.lock()) {
        s->StartImpl(std::move(yc));
      } else {
        LOG(ERROR) << "scheduler is already destroyed";
      }
    });
  } else {
    THROW() << "scheduler has no context!";
  }
}
